package com.googlewell.zelix.model;

import com.googlewell.zelix.tool.ClassTool;
import com.googlewell.zelix.tool.SearchTool;
import me.coley.recaf.control.Controller;
import me.coley.recaf.search.MemberResult;
import me.coley.recaf.search.SearchCollector;
import me.coley.recaf.search.SearchResult;
import me.coley.recaf.search.StringMatchMode;
import me.coley.recaf.util.ClassUtil;
import org.jetbrains.java.decompiler.struct.gen.MethodDescriptor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.tree.InvokeDynamicInsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;

import java.util.ArrayList;
import java.util.HashSet;

public class Zelix {
    /**
     * 类名
     */
    private String className = "com/zelix/ZKM";
    /**
     * 主方法名
     */
    private String memberName = "main";
    /**
     * 主方法描述
     */
    private String memberDesc = "([Ljava/lang/String;)V";

    /**
     * Recaf Controller
     */
    private Controller controller;
    /**
     * main方法
     */
    private MethodNode mainMethod;
    /**
     * java字节码汇编指令
     */
    private String[] codeLines;
    /**
     * 内部InvokeDynamic,多用于解密字符串
     */
    private InvokeDynamicInsnNode innerInvokeDynamic;
    /**
     * 外部InvokeDynamic,多用于成员属性和方法调用
     */
    private InvokeDynamicInsnNode outerInvokeDynamic;
    /**
     * 计算long值的接口类
     */
    private MemberResult interfaceResult;
    /**
     * 计算long值的接口类的实现类
     */
    private HashSet<String> classReferences;

    public Zelix(Controller controller) {
        this.controller = controller;
        setMainMethod();
        setCodeLines();
        setInnerInvokeDynamic();
        setOuterInvokeDynamic();
        setInterfaceResult();
        setClassReferences();
    }

    private void setMainMethod() {
        //获取类的ClassReader
        ClassReader cr = controller.getWorkspace().getClassReader(className);
        //获取main方法MethodNode
        if (cr == null) {
            mainMethod = null;
        } else {
            mainMethod = ClassUtil.getMethod(cr, 0, memberName, memberDesc);
        }
    }

    private void setCodeLines() {
        if (mainMethod == null) {
            codeLines = null;
        } else {
            codeLines = ClassTool.getMethodCodeLines(mainMethod);
        }
    }

    private void setInnerInvokeDynamic() {
        innerInvokeDynamic = ClassTool.searchInvokeDynamic(className, codeLines, true);
    }

    private void setOuterInvokeDynamic() {
        outerInvokeDynamic = ClassTool.searchInvokeDynamic(className, codeLines, false);
    }

    private void setInterfaceResult() {
        if (codeLines == null) {
            interfaceResult = null;
            return;
        }
        //分析long值计算的类
        //查找指令所在行
        String ins = "INVOKEVIRTUAL java/lang/invoke/MethodHandles$Lookup.lookupClass()Ljava/lang/Class;";
        int line = ClassTool.getCodeLine(codeLines, ins);
        //不存在下一条指令
        if (line == -1 || line + 1 >= codeLines.length) {
            interfaceResult = null;
            return;
        }

        //获取下一条MethodInsnNode指令
        MethodInsnNode methodInsnNode = ClassTool.getMethodInsnNodeByInsnString(codeLines[line + 1]);
        //不是method指令
        if (methodInsnNode == null) {
            interfaceResult = null;
            return;
        }
        //以上搜索指令和MethodInsnNode可考虑循环读取符合条件的,知道搜索完所有指令,而不是只搜索开头,按现在情况暂时够用了
        //System.out.println(methodInsnNode.getOpcode() + ":" + OpcodeUtil.opcodeToName(methodInsnNode.getOpcode()) + ":" + methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc);

        //获取方法描述
        MethodDescriptor methodDescriptor = MethodDescriptor.parseDescriptor(methodInsnNode.desc);
        //获取方法返回类型,比如返回类型描述Lcom/zelix/cmm;  methodDescriptor.ret.value为com/zelix/cmm    methodDescriptor.ret.toString()为Lcom/zelix/cmm;
        //System.out.println("接口类:" + methodDescriptor.ret.value);

        //搜索这个接口类的成员方法,以便获取方法名
        SearchCollector intfMemColl = SearchTool.searchMemberDefinition(controller.getWorkspace(), methodDescriptor.ret.value, null, "(J)J", StringMatchMode.EQUALS, new ArrayList<>());
        if (intfMemColl.getAllResults().size() == 1) {//只有一个
            if (intfMemColl.getAllResults().get(0) instanceof MemberResult) {
                MemberResult memberResult = (MemberResult) intfMemColl.getAllResults().get(0);
                if (memberResult.isMethod()) {
                    //System.out.println("接口类:" + memberResult.getOwner() + ":" + memberResult.getName() + memberResult.getDesc());
                    interfaceResult = memberResult;
                    return;
                }
            }
        }
        interfaceResult = null;
    }

    private void setClassReferences() {
        //搜索这个类的引用
        HashSet<String> classs = new HashSet<>();
        if (interfaceResult != null) {
            SearchCollector clsRefColl = SearchTool.searchClassReference(controller.getWorkspace(), interfaceResult.getOwner(), StringMatchMode.CONTAINS, new ArrayList<>());
            for (SearchResult result : clsRefColl.getAllResults()) {
                classs.add(ClassTool.getClassContext(result.getContext()).getName());
            }
        }
        classReferences = classs;
    }

    public InvokeDynamicInsnNode getInnerInvokeDynamic() {
        return innerInvokeDynamic;
    }

    public InvokeDynamicInsnNode getOuterInvokeDynamic() {
        return outerInvokeDynamic;
    }

    public MemberResult getInterfaceResult() {
        return interfaceResult;
    }

    public HashSet<String> getClassReferences() {
        return classReferences;
    }
}
